#!/usr/bin/env python3
"""
数据库迁移脚本 - 添加缺失的字段
"""

import os
import sys
sys.path.insert(0, os.path.dirname(os.path.dirname(__file__)))

from flask import Flask
from config import get_config
from src.models.database import db
import pymysql

def migrate_database():
    """执行数据库迁移"""
    print("🔄 开始数据库迁移...")
    
    # 创建Flask应用
    app = Flask(__name__)
    config = get_config()
    app.config.from_object(config)
    
    try:
        # 初始化数据库
        db.init_app(app)
        
        with app.app_context():
            # 获取数据库连接
            connection = db.engine.raw_connection()
            cursor = connection.cursor()
            
            print("📊 检查表结构...")
            
            # 检查sector_id字段是否存在
            cursor.execute("""
                SELECT COUNT(*) FROM information_schema.COLUMNS 
                WHERE table_name = 'metaso_query_data' 
                AND table_schema = 'qilu_db' 
                AND column_name = 'sector_id'
            """)
            
            sector_id_exists = cursor.fetchone()[0] > 0
            
            # 检查sector_name字段是否存在
            cursor.execute("""
                SELECT COUNT(*) FROM information_schema.COLUMNS 
                WHERE table_name = 'metaso_query_data' 
                AND table_schema = 'qilu_db' 
                AND column_name = 'sector_name'
            """)
            
            sector_name_exists = cursor.fetchone()[0] > 0
            
            # 添加缺失的字段
            if not sector_id_exists:
                print("➕ 添加sector_id字段...")
                cursor.execute("""
                    ALTER TABLE metaso_query_data 
                    ADD COLUMN sector_id VARCHAR(50) COMMENT '板块ID'
                """)
                print("✅ sector_id字段添加成功")
            else:
                print("ℹ️  sector_id字段已存在")
            
            if not sector_name_exists:
                print("➕ 添加sector_name字段...")
                cursor.execute("""
                    ALTER TABLE metaso_query_data 
                    ADD COLUMN sector_name VARCHAR(100) COMMENT '板块名称'
                """)
                print("✅ sector_name字段添加成功")
            else:
                print("ℹ️  sector_name字段已存在")
            
            # 提交更改
            connection.commit()
            
            # 显示最终表结构
            print("\n📋 最终表结构:")
            cursor.execute("DESCRIBE metaso_query_data")
            columns = cursor.fetchall()
            
            for column in columns:
                print(f"  - {column[0]}: {column[1]} {column[2] if column[2] == 'YES' else 'NOT NULL'}")
            
            cursor.close()
            connection.close()
            
            print("\n🎉 数据库迁移完成!")
            return True
            
    except Exception as e:
        print(f"❌ 数据库迁移失败: {str(e)}")
        return False

if __name__ == "__main__":
    print("🔧 数据库迁移工具")
    print("=" * 50)
    
    success = migrate_database()
    
    if success:
        print("\n✅ 迁移成功，现在可以启动应用!")
        print("运行命令: python src/main.py")
    else:
        print("\n❌ 迁移失败，请检查错误信息")
